rm(list=ls())

# Load packages
library(foreign)
library(lmtest)
library(sandwich)
library(stargazer)
library(msm)
library(readxl)
library(ggplot2)
library(gridExtra)
library(readstata13)
library(haven)
library(dplyr)
library(forcats)
library(ggeasy)

# Load data
cep_all <- read_dta("Data/cep_all_other_outcomes.dta")

# Define dosage range samples
dosage_ranges <- c(1, 2, 3, 4, 5, 6)
sample_list <- lapply(dosage_ranges, function(i) subset(cep_all, treatment_dosage_coup >= -i & treatment_dosage_coup <= i))
names(sample_list) <- paste0("sample", dosage_ranges)

# Clustered SE function
cluster_se <- function(model, data, type = "HC1") {
  idx <- as.integer(rownames(model.frame(model)))
  vc  <- sandwich::vcovCL(model, cluster = data$cluster[idx], type = type)
  lmtest::coeftest(model, vcov. = vc)
}

# Reusable model block generator
run_models <- function(samples, formula, dv_name, table_title, dep_caption, column_labels, cov_label) {
  models_lm <- lapply(samples, function(s) lm(formula, data = s))
  models_cl <- Map(cluster_se, models_lm, samples)
  sample_sizes <- sapply(models_lm, nobs)
  r_squared <- sapply(models_lm, function(m) summary(m)$r.squared)
  mean_dv <- sapply(models_lm, function(m) mean(m$model[[dv_name]], na.rm = TRUE))
  
  stargazer(models_cl,
            type = "latex",
            title = table_title,
            keep = "^treat1$",
            column.labels = column_labels,
            covariate.labels = cov_label,
            dep.var.caption = dep_caption,
            dep.var.labels = dv_name,
            dep.var.labels.include = FALSE,
            add.lines = list(
              c("Obs.", sample_sizes),
              c("$R^2$", round(r_squared, 3)),
              c("Mean DV", round(mean_dv, 3))
            ),
            align = TRUE)
}

# ------------------
# Table C1: Democracy Works Poorly
table_title_c1 <- "Table C1: RD Estimates 1973 Coup on Views About How Democracy Works"
dep_caption_c1 <- "Dependent Variable: Democracy Works Poorly"

tabc1 <- run_models(sample_list[1:5],
           democracy_bad ~ treat1 + treatment_dosage_coup + treatment_dosage_coup:treat1 + as.factor(encuesta),
           dv_name = "democracy_bad",
           table_title = table_title_c1,
           dep_caption = dep_caption_c1,
           column_labels = c("54 and 56", "53-57", "52-58", "51-59", "50-60"),
           cov_label = "Coup")

print('Table C1 complete')

# ------------------
# Table C2: Democracy Always Preferable
table_title_c2 <- "Table C2: RD Estimates 1973 Coup on Normative Views About Democracy"
dep_caption_c2 <- "Dependent Variable: Democracy Always Preferable"

tabc2 <- run_models(sample_list[1:5],
           dem_3 ~ treat1 + treatment_dosage_coup + treatment_dosage_coup:treat1 + as.factor(encuesta),
           dv_name = "dem_3",
           table_title = table_title_c2,
           dep_caption = dep_caption_c2,
           column_labels = c("54 and 56", "53-57", "52-58", "51-59", "50-60"),
           cov_label = "Coup")

print('Table C2 complete')